#ifndef UHM_DISCRETE_EXTENSION_MESH2TREE_HXX_
#define UHM_DISCRETE_EXTENSION_MESH2TREE_HXX_

// METIS
extern "C" {

#include "defs.h"
#include "struct.h"
#include "macros.h"
#include "rename.h"

  extern void SetUpGraph(GraphType *, Int_, Int_, Int_,
                         idxtype *, idxtype *, idxtype *, idxtype *, Int_);
  extern Int_ idxsum(Int_, idxtype *);
  extern void InitRandom(Int_);
  extern void AllocateWorkSpace(CtrlType *, GraphType *, Int_);
  extern void FreeWorkSpace(CtrlType *, GraphType *);
  extern void MlevelEdgeBisection(CtrlType *, GraphType *, Int_ *, float);
  extern void SplitGraphPart(CtrlType *,
                             GraphType *, GraphType *, GraphType *);
  extern void GKfree(void **,...);
}


namespace uhm {

  template<class M_, class T_> 
  Bool_ mesh2tree(M_ &mesh, T_ &tree, Int_ pid=0);

  template<class M_, class T_>
  Bool_ mesh2tree_recursive(std::vector< typename M_::Element_* > *els,
                            M_ *mesh, T_ *tree, typename T_::Tnd_* p,
                            CtrlType *ctrl, GraphType *graph, Int_ parallel);
  
  // organize tree for given partition, pid
  template<class M_, class T_> 
  Bool_ mesh2tree(M_ &mesh, T_ &tree, Int_ pid) {

    std::vector< typename M_::Element_* > els;
    mesh.dump(els, pid);

    time_in("uhm:: mesh2tree",UHM_TIME_LEVEL_INTERNAL);

    if (els.size() > 1) {

      Int_ level=0, n=els.size(), wgtflag=1;
      for (Int_ i=n;i>0;i/=2)
        ++level;

      // MUST DO... set size n*log2(n)
      // reallocation needs to create all object again and 
      // those are not same as before. I create object with shallow copying...
      tree.clear();
      tree.reserve(n*level);

      // Root
      tree.set_root(tree.add_node(mesh.add_element()));

      std::vector< Int_ > xadj, adjncy, adjwgt;
      elements2graph<M_>(&els, &xadj, &adjncy, &adjwgt, pid);

      CtrlType ctrl;
      GraphType graph;

      // Create METIS graph
      SetUpGraph(&graph, OP_OEMETIS, n, 1, &xadj[0], &adjncy[0],
                 NULL, &adjwgt[0], wgtflag);

      // Default METIS parameter
      ctrl.CType   = PMETIS_CTYPE;
      ctrl.IType   = PMETIS_ITYPE;
      ctrl.RType   = PMETIS_RTYPE;
      ctrl.dbglvl  = PMETIS_DBGLVL;

      ctrl.optype    = OP_PMETIS;
      ctrl.CoarsenTo = 20;
      ctrl.maxvwgt   = 1.5*(idxsum(n, graph.vwgt)/ctrl.CoarsenTo);

      InitRandom(-1);


      time_in("uhm:: mesh2tree:: Recursion",UHM_TIME_LEVEL_INTERNAL);
#pragma omp parallel
      {
#pragma omp single nowait
        {
          AllocateWorkSpace(&ctrl, &graph, 2);
          mesh2tree_recursive<M_,T_>(&els, &mesh, &tree,
                                     tree.get_root(),  
                                     &ctrl, &graph, true);

          FreeWorkSpace(&ctrl, &graph);
        }
      }
      time_out(UHM_TIME_LEVEL_INTERNAL);

    } else if (els.size() == 1) {
      Int_ n=els.size();
      tree.clear();
      tree.reserve(n*2);
      tree.set_root(tree.add_node(els[0]));
    }
    time_out(UHM_TIME_LEVEL_INTERNAL);

      time_in("uhm:: mesh2tree:: Sorting Tree",UHM_TIME_LEVEL_INTERNAL);

	  tree.set_gen();
    tree.sort();
	  time_out(UHM_TIME_LEVEL_INTERNAL);

    return true;
  }

  template<class M_, class T_>
  Bool_ mesh2tree_recursive(std::vector< typename M_::Element_* > *els,
                            M_ *mesh, T_ *tree, typename T_::Tnd_* p,
                            CtrlType *ctrl, GraphType *graph, Int_ parallel) {

    Int_ n=graph->nvtxs, tvwgt, tpwgts2[2];
    CtrlType  ctrl_self[2];
    GraphType graph_self[2];

    tvwgt = idxsum(n, graph->vwgt);
    tpwgts2[0] = tvwgt/2;
    tpwgts2[1] = tvwgt-tpwgts2[0];

    MlevelEdgeBisection(ctrl, graph, tpwgts2, 1.0);
    SplitGraphPart(ctrl, graph, &graph_self[0], &graph_self[1]);

    for (Int_ i=0;i<2;++i) {
      if (graph_self[i].nvtxs > 3) {
        typename T_::Tnd_* t = 
          tree->add_node(mesh->add_element());
        t->set_parent(p);
        p->add_child(t);

        // Inherit same control parameters
        ctrl_self[i] = *ctrl;

        if (graph_self[i].nvtxs > UHM_NESTED_DISSECTION_TASK_SIZE && parallel) {

#pragma omp task firstprivate(els, mesh, tree, i, t, ctrl_self, graph_self)
          {
            AllocateWorkSpace(&ctrl_self[i], &graph_self[i], 2);
            mesh2tree_recursive(els, mesh, tree, t,
                                &ctrl_self[i], &graph_self[i], true);
            FreeWorkSpace(&ctrl_self[i], &graph_self[i]);
          }

        } else {
          mesh2tree_recursive(els, mesh, tree, t,
                              &ctrl_self[i], &graph_self[i], false);
        }

      } else {

        if ( graph_self[i].nvtxs ) {
          typename T_::Tnd_* t;
          switch ( graph_self[i].nvtxs ) {
          case 1:
            t = p;
            break;
          default:
            t = tree->add_node(mesh->add_element());
            t->set_parent(p);
            p->add_child(t);
            break;
          }

          for (Int_ j=0;j<graph_self[i].nvtxs;++j) {
            typename T_::Tnd_ *c = 
              tree->add_node(els->at(graph_self[i].label[j]));
            c->set_parent(t);
            t->add_child(c);
          }

          GKfree((void**)(&graph_self[i].gdata),
                 (void**)(&graph_self[i].rdata),
                 (void**)(&graph_self[i].label), LTERM);
        }
      }
    }
    GKfree((void**)(&graph->gdata),
           (void**)(&graph->rdata),
           (void**)(&graph->label), LTERM);

    return true;
  }
}

#endif
